"""
用于打印gpu状态的工具类
"""


import time
import multiprocessing
import subprocess
import psutil


print_event = multiprocessing.Event()
# 创建一个事件，用于通知子进程何时停止
stop_event = multiprocessing.Event()


# 打印 python 版本信息
def print_python_version():
    import sys
    print(f"Python version: {sys.version}")


def is_cuda_available():
    try:
        print_python_version()

        import torch
        # 检查 CUDA 是否可用
        print(f"Is CUDA available: {torch.cuda.is_available()}")
        # True
        print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
        # Tesla T4
    except Exception as e:
        print(f"cuda 不可用")
        print(f"Error: {e}")


def print_gpu():
    # 打印 GPU 信息, 需要系统有 ixsmi 命令
    try:
        result = subprocess.run(['ixsmi'], capture_output=True, text=True)
        print('ixsmi Output:', result.stdout)
    except Exception as e:
        print('没有 ixsmi 命令')


def print_cpu():
    # 打印CPU信息
    print("CPU信息:")
    cpu_count = psutil.cpu_count(logical=False)  # 获取物理CPU数量
    print(f"物理CPU数量: {cpu_count}")
    cpu_count_logical = psutil.cpu_count()  # 获取逻辑CPU数量
    print(f"逻辑CPU数量: {cpu_count_logical}")
    cpu_freq = psutil.cpu_freq()  # 获取CPU频率
    print(f"CPU频率: {cpu_freq.current} MHz")

    # 打印内存信息
    print("\n内存信息:")
    memory = psutil.virtual_memory()
    print(f"总内存: {memory.total / (1024 ** 3):.2f} GB")
    print(f"可用内存: {memory.available / (1024 ** 3):.2f} GB")
    print(f"已使用内存: {memory.used / (1024 ** 3):.2f} GB")
    print(f"内存使用率: {memory.percent}%")

    # 打印当前CPU使用率
    print("\n当前CPU使用率:")
    cpu_usage = psutil.cpu_percent(interval=1)  # interval参数表示间隔时间，单位为秒
    print(f"CPU使用率: {cpu_usage}%")

    # 打印当前内存使用率
    print("\n当前内存使用率:")
    memory_usage = memory.percent
    print(f"内存使用率: {memory_usage}%")


def print_cpu_oneline():
    cpu_count = psutil.cpu_count(logical=False)  # 获取物理CPU数量
    cpu_count_logical = psutil.cpu_count()  # 获取逻辑CPU数量
    cpu_freq = psutil.cpu_freq()  # 获取CPU频率
    cpu_usage = psutil.cpu_percent(interval=1)  # interval参数表示间隔时间，单位为秒
    print(f"物理CPU数量: {cpu_count} 逻辑CPU数量: {cpu_count_logical} CPU频率: {cpu_freq.current} MHz  CPU使用率: {cpu_usage}%")

    memory = psutil.virtual_memory()
    print(f"总内存: {memory.total / (1024 ** 3):.2f} GB  可用内存: {memory.available / (1024 ** 3):.2f} GB  已使用内存: {memory.used / (1024 ** 3):.2f} GB  内存使用率: {memory.percent}%")


def print_gpu_process(event, stop_event):
    while not stop_event.is_set():  # 检查事件是否被设置
        if event.is_set():
            print_gpu()
            # print_cpu()
            print_cpu_oneline()
        time.sleep(2)
    print("子进程结束")


def start_print_gpu_process(event=print_event, stop_event=stop_event):
    # 创建一个新的进程
    process = multiprocessing.Process(target=print_gpu_process, args=(event, stop_event,))
    # 启动子进程
    process.start()
    print("启动子进程")
    return process




def print_set():
    global print_event
    print_event.set()


def print_unset():
    global print_event
    print_event.clear()

def stop_set():
    global stop_event
    stop_event.set()

def stop_unset():
    global stop_event
    stop_event.clear()




